import torch
import torch.nn as nn
import time
import os
from multiprocessing import Process, Lock

"""
cmd_file: XXX.cmd
content:
    # gres=gpu:1
    pushd YYY
    rm ./KKK.log
    python KKK.py >> KKK.log
"""


devices = [str(device_id) for device_id in range(torch.cuda.device_count())]
runing_devices = []
taskNames = []
monitor_Dir = "./{}".format(str(__file__).rstrip(".py"))
completed_dir = "{}_completed_tasks/".format(monitor_Dir)

def runCommandFile(cmd_file, lock):
    this_devices = []
    deviceCount = 0
    lines = []
    old_cmd_file = "{}/{}".format(monitor_Dir, cmd_file)
    with open(old_cmd_file, 'r') as fr:
        for line in fr:
            if "gres=gpu:" in line:
                deviceCount = int(line.split("gres=gpu:")[1])
                lock.acquire()
                if deviceCount > len(devices):
                    print("Error!!! Job {} lacks the available devices!".format(cmd_file))
                    lock.release()
                    break
                this_devices = [devices.pop(0) for _ in range(deviceCount)]
                runing_devices.append(this_devices)
                lock.release()
                lines.append(line)
            elif ("python " in line) and (".py" in line):
                assert len(this_devices) == deviceCount
                lines.append("CUDA_VISIBLE_DEVICES={} python {}".format(
                    ','.join(this_devices), line.split("python ")[1]
                ))
            else:
                lines.append(line)

    new_cmd_file = "{}/{}".format(monitor_Dir, cmd_file.replace(".cmd", ".sh"))
    with open(new_cmd_file, 'w') as fw:
        fw.write("".join(lines))
    print("cmd_file : ", cmd_file)
    print("old_cmd_file : ", old_cmd_file)
    print("new_cmd_file : ", new_cmd_file)
    print(lines)
    print("task {} started! ".format(cmd_file))
    print(os.system("bash {}".format(new_cmd_file)))
    print("task {} is completed!".format(cmd_file))

    lock.acquire()
    devices.extend(this_devices)
    os.system("mv {}/{} ./{}/{}".format(
        monitor_Dir, cmd_file, completed_dir, cmd_file)
    )
    os.system("mv {} {}/{}".format(
        new_cmd_file, completed_dir, new_cmd_file)
    )
    lock.release()

def runRoutineTask(lock):
    lock.acquire()
    # print("+++run RoutineTask++++ ")
    os.system("CUDA_VISIBLE_DEVICES={} python ./{}/{}.py".format(','.join(devices), monitor_Dir, monitor_Dir))
    # print("+++Complete RoutineTask++++ ")
    lock.release()

if __name__ == "__main__":
    print("main start")
    tasks = []
    if not os.path.exists(monitor_Dir):
        os.system("mkdir {}".format(monitor_Dir))
    if not os.path.exists(completed_dir):
        os.system("mkdir {}".format(completed_dir))
    l = Lock()
    assert torch.cuda.device_count() > 1
    while True:
        l.acquire()
        l.release()
        cmd_files = [fname for fname in os.listdir(monitor_Dir) if ".cmd" in fname]
        # if the file is not in the dir now, the task is completed
        taskNames = [fname for fname in taskNames if fname in cmd_files] # taskNames = taskNames & cmd_files
        for fname in cmd_files:
            if fname in taskNames:
                pass
            else:
                p = Process(target=runCommandFile, args=(fname, l, ))
                p.start()
                taskNames.append(fname)
        main_p = Process(target=runRoutineTask, args=(l, ))
        main_p.start()
        time.sleep(2)
    print("main complete")